# -*- coding: utf-8 -*-
# @Time    : 2023/5/20 15:32
# @Author  : Pan
# @Software: PyCharm
# @Project : VisualFramework
# @FileName: MAE

import os
import cv2
import time
import loss
import paddle
import networks
import optimizers
import datasets
import numpy as np
from core import base
from visualdl import LogWriter


class MAEEngine:
    def __init__(self, config):
        self.start_time = time.time()
        self.next_time = time.time()
        self.now_time = time.time()

        self.base_info_config = config["base_info"]
        self.train_dataset_config = config["train_dataset"]
        self.amp = config["amp"] if "amp" in config.keys() else None
        self.optimizer_config = config["optimizer"]
        self.network_config = config["network"]
        self.loss_config = config["loss"]

        self.writer = LogWriter(logdir=self.base_info_config["log_dir"])

        self.model = networks.make_model(self.network_config)
        self.model, self.optimizer, self.lr = optimizers.make_optim(self.optimizer_config, self.model)
        self.train_dataloader = datasets.make_dataloader(self.train_dataset_config)
        self.loss = loss.LossCompose(self.loss_config)

        if self.base_info_config["pretrained"] is None:
            self.step = 0
        else:
            self.model, self.optimizer, self.lr, self.step = base.load_model(self.model, self.optimizer, self.lr, self.base_info_config["pretrained"])

        if self.amp:
            self.amp = paddle.amp.GradScaler(init_loss_scaling=self.amp["scale"])

    def train(self):
        if self.amp is not None:
            self.fp16_train()
        else:
            while self.step < self.base_info_config["step"]:
                for idx, data in enumerate(self.train_dataloader):
                    self.step += 1
                    img, mask = data["img"], data["mask"]
                    self.optimizer.clear_grad()
                    predict = self.model(img, mask)
                    loss_value = self.loss(predict, [img], mask=mask)
                    loss_value.backward()
                    self.optimizer.step()
                    self.lr.step()
                    if self.step % self.base_info_config["dot"] == 0:
                        self.display()
                    if self.step % self.base_info_config["save_iters"] == 0:
                        print("\033[5;31;47m[display]\033[0m save iters: %8d in %s" % ( self.step, self.base_info_config["save_path"]))
                        base.save_model(self.model, self.optimizer, self.step, self.base_info_config["save_path"], only_last=True)
                        self.draw_image(data, predict[0], mask)

    def draw_image(self, data, predicts, mask):
        for i in range(4 if predicts.shape[0] > 4 else predicts.shape[0]):
            predict = np.clip(predicts[i].transpose([1, 2, 0]).numpy() * 127.5 + 127.5, 0, 255).astype(np.uint8)
            image = np.clip(data["img"][i].transpose([1, 2, 0]).numpy() * 127.5 + 127.5, 0, 255).astype(np.uint8)
            mask_ = np.tile(1 - mask[i], (3, 1, 1)).transpose([1, 2, 0]).astype(np.uint8)
            mask_ = cv2.resize(mask_, image.shape[:2], interpolation=cv2.INTER_NEAREST)
            input_image = (mask_ * image).astype(np.uint8)
            generate_image = (mask_ * image).astype(np.uint8) + ((1-mask_) * predict).astype(np.uint8)

            cv2.putText(input_image, "input", (10, 10), cv2.FONT_HERSHEY_SIMPLEX, 0.75, (0, 0, 255), 1)
            cv2.putText(image, "target", (10, 10), cv2.FONT_HERSHEY_SIMPLEX, 0.75, (0, 0, 255), 1)
            cv2.putText(predict, "predict", (10, 10), cv2.FONT_HERSHEY_SIMPLEX, 0.75, (0, 0, 255), 1)
            cv2.putText(generate_image, "generate", (10, 10), cv2.FONT_HERSHEY_SIMPLEX, 0.75, (0, 0, 255), 1)
            image = np.concatenate([image, predict, input_image, generate_image], axis=1)
            cv2.imwrite(os.path.join(self.base_info_config["save_path"], str(self.step), os.path.basename(data["path"][i])), image)

    def display(self):
        self.now_time = time.time()
        epoch = int(self.step / len(self.train_dataloader))
        process = self.step/self.base_info_config["step"]
        speed_time = self.now_time - self.start_time
        remain_time = (self.now_time - self.next_time) / (self.base_info_config["dot"] / self.base_info_config["step"]) * (1 - process)
        self.next_time = self.now_time
        info_list = [
            {
                "name": "learning_rate",
                "value": self.lr.get_lr()
            }
        ] + self.loss.get_loss_info()
        base_info = "\033[0;31;40m[Train]\033[0m %s epochs:%4d steps:%9d/%9d process:%5.2f%% speed_time:%s remain_time:%s" % (
            time.ctime(), epoch, self.step, self.base_info_config["step"], process*100, base.time_std(speed_time), base.time_std(remain_time)
        )
        for item in info_list:
            self.writer.add_scalar(tag="train/" + item["name"], step=self.step, value=item["value"])
            base_info += " %s:%f" % (item["name"], item["value"])
        print(base_info)

    def fp16_train(self):
        while self.step < self.base_info_config["step"]:
            for idx, data in enumerate(self.train_dataloader):
                self.step += 1
                img, mask = data["img"], data["mask"]
                self.optimizer.clear_grad(set_to_zero=False)
                with paddle.amp.auto_cast(custom_white_list={'elementwise_add'}, level='O1'):
                    predict = self.model(img, mask)
                    loss_value = self.loss(predict, [img], mask=mask)
                scaled = self.amp.scale(loss_value)
                scaled.backward()
                self.amp.step(self.optimizer)
                self.amp.update()
                self.lr.step()
                if self.step % self.base_info_config["dot"] == 0:
                    self.display()
                if self.step % self.base_info_config["save_iters"] == 0:
                    print("save iters: %8d in %s" % (self.step, self.base_info_config["save_path"]))
                    base.save_model(self.model, self.optimizer, self.step, self.base_info_config["save_path"], only_last=True)
                    self.draw_image(data, predict[0], mask)
